import chromadb
from sentence_transformers import SentenceTransformer
import numpy as np
from typing import List, Dict
import os

class VectorStore:
    """向量存储和检索系统"""
    
    def __init__(self, persist_directory: str = "chroma_db"):
        self.persist_directory = persist_directory
        self.collection_name = "documents"
        
        # 初始化向量化模型
        print("正在加载向量化模型...")
        self.embedding_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
        print("向量化模型加载完成")
        
        # 初始化ChromaDB
        self.chroma_client = chromadb.PersistentClient(path=persist_directory)
        
        # 获取或创建集合
        try:
            self.collection = self.chroma_client.get_collection(name=self.collection_name)
        except Exception:
            self.collection = self.chroma_client.create_collection(name=self.collection_name)
    
    def add_documents(self, documents: List[Dict[str, str]], filename: str):
        """添加文档到向量存储"""
        if not documents:
            return
        
        # 准备数据
        texts = [doc["content"] for doc in documents]
        ids = [f"{filename}_{doc['chunk_id']}" for doc in documents]
        metadatas = [{"filename": filename, "chunk_id": doc["chunk_id"]} for doc in documents]
        
        # 生成嵌入向量
        embeddings = self.embedding_model.encode(texts).tolist()
        
        # 添加到集合
        self.collection.add(
            embeddings=embeddings,
            documents=texts,
            metadatas=metadatas,
            ids=ids
        )
        
        print(f"成功添加 {len(documents)} 个文档块")
    
    def search(self, query: str, k: int = 3) -> List[Dict]:
        """搜索相关文档"""
        # 生成查询向量
        query_embedding = self.embedding_model.encode([query]).tolist()
        
        # 执行搜索
        results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=k
        )
        
        # 格式化结果
        formatted_results = []
        if results["documents"] and results["documents"][0]:
            for i in range(len(results["documents"][0])):
                formatted_results.append({
                    "content": results["documents"][0][i],
                    "metadata": results["metadatas"][0][i],
                    "score": 1 - results["distances"][0][i]  # 转换为相似度分数
                })
        
        return formatted_results
    
    def get_document_count(self) -> int:
        """获取文档数量"""
        try:
            return self.collection.count()
        except Exception:
            return 0
    
    def clear(self):
        """清除所有文档"""
        # 删除集合并重新创建
        try:
            self.chroma_client.delete_collection(name=self.collection_name)
            self.collection = self.chroma_client.create_collection(name=self.collection_name)
        except Exception as e:
            print(f"清除文档失败: {e}")
